package com.tutu.mybatis.dynamic.datasource.config;

import com.tutu.mybatis.dynamic.datasource.handler.CustomizeHandlerMethodReturnValueHandler;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.method.support.HandlerMethodReturnValueHandler;
import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerAdapter;
import org.springframework.web.servlet.mvc.method.annotation.RequestResponseBodyMethodProcessor;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

/**
 * @author qyy
 * @date 2021/3/19 17:52
 */
@Configuration
public class ReturnValueConfig implements InitializingBean {

    @Resource
    RequestMappingHandlerAdapter requestMappingHandlerAdapter;

    @Override
    public void afterPropertiesSet() throws Exception {
        List<HandlerMethodReturnValueHandler> originHandlers = requestMappingHandlerAdapter.getReturnValueHandlers();
        if (originHandlers != null) {
            List<HandlerMethodReturnValueHandler> newHandlers = new ArrayList<>(originHandlers.size());
            for (HandlerMethodReturnValueHandler originHandler : originHandlers) {
                if (originHandler instanceof RequestResponseBodyMethodProcessor) {
                    newHandlers.add(new CustomizeHandlerMethodReturnValueHandler(originHandler));
                }else{
                    newHandlers.add(originHandler);
                }
            }
            requestMappingHandlerAdapter.setReturnValueHandlers(newHandlers);
        }
    }
}
